import pygame, random, sys, os, time, pickle, tkinter, tkinter.filedialog, tkinter.ttk
import pygame.gfxdraw
from simulation_engine import SimulationEngine, Cell
from cell_genetics import CellGenetics
from ui_components import Panel, Text, GeneStatDisplay


class BioSimulator:
    def __init__(self):
        pygame.init()
        pygame.font.init()
        self.screen = pygame.display.set_mode((1200, 800))
        self.fullscreen = False
        self.clock = pygame.time.Clock()
        self.sim = SimulationEngine(120, 80)
        self.start_time = None
        # 初始化一些随机细胞
        for _ in range(100):
            x, y = random.randint(0, 79), random.randint(0, 59)
            if x < len(self.sim.grid) and y < len(self.sim.grid[0]):
                self.sim.grid[x][y] = Cell()
        self.selected_cell = None
        self.selected_genetics = None
        self.running = False
        self.grid_size = 10
        self.panel_rect = pygame.Rect(900, 0, 300, 800)
        self.control_rect = pygame.Rect(10, 10, 150, 40)
        self.panel = Panel(self.panel_rect)
        self.cell_count_text = Text("Cells: 0", (910, 10), (255, 255, 255))
        self.genetic_info_text = Text("Gene: ", (910, 50), (255, 255, 255))
        self.gene_stat_font = pygame.font.SysFont('Arial', 16)
        self.gene_stat_rect = pygame.Rect(910, 100, 280, 300)
        self.gene_stat_display = GeneStatDisplay(self.gene_stat_rect, self.gene_stat_font)
        self.panel.add_component(self.gene_stat_display)
        self.total_generations_text = Text("Total Generations: 0", (910, 90), (255, 255, 255))

        # 添加功能按钮
        self.save_button = Text("Save", (910, 650), (255, 255, 255))
        self.load_button = Text("Load", (910, 680), (255, 255, 255))
        self.new_button = Text("New", (910, 710), (255, 255, 255))
        self.analyze_button = Text("Analyze", (910, 740), (255, 255, 255))
        self.random_button = Text("Random", (910, 770), (255, 255, 255))

    def draw_grid(self):
        for x in range(len(self.sim.grid)):
            for y in range(len(self.sim.grid[0])):
                if self.sim.grid[x][y]:
                    color = self.get_cell_color(self.sim.grid[x][y])
                    if color:
                        pygame.draw.rect(self.screen, color,
                                       (x * self.grid_size, y * self.grid_size, self.grid_size - 1, self.grid_size - 1))
                    if self.selected_genetics and self.sim.grid[x][y].genetics == self.selected_genetics:
                        pygame.gfxdraw.rectangle(self.screen, (x * self.grid_size, y * self.grid_size, self.grid_size, self.grid_size), (255, 255, 255))

    def get_cell_color(self, cell):
        genes = cell.genetics.decode_gene()
        # 使用基因哈希值作为色相基础
        gene_hash = hash(tuple(cell.genetics.genetic_code))
        hue = gene_hash % 360  # 色相范围0-360

        # 根据细胞年龄调整饱和度
        saturation = max(50, 100 - min(cell.age / 10, 100))

        # 根据能量调整亮度
        value = max(50, min(cell.energy / 2, 100))

        # 转换HSV到RGB
        h = hue / 60
        c = value * saturation / 10000
        x = c * (1 - abs(h % 2 - 1))

        if h < 1:
            r, g, b = c, x, 0
        elif h < 2:
            r, g, b = x, c, 0
        elif h < 3:
            r, g, b = 0, c, x
        elif h < 4:
            r, g, b = 0, x, c
        elif h < 5:
            r, g, b = x, 0, c
        else:
            r, g, b = c, 0, x

        m = value / 100 - c
        r, g, b = (r + m) * 255, (g + m) * 255, (b + m) * 255

        # 确保颜色值在有效范围内
        r = max(0, min(255, int(r)))
        g = max(0, min(255, int(g)))
        b = max(0, min(255, int(b)))

        return (r, g, b)

    def run(self):
        while True:
            # 主循环
            # 事件处理
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    sys.exit()
                if event.type == pygame.MOUSEBUTTONDOWN:
                    self.handle_click(event.pos)
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_SPACE:
                        self.running = not self.running
                if event.type == pygame.KEYDOWN and event.key == pygame.K_F11:
                    self.toggle_fullscreen()

            # 模拟更新
            if self.running:
                self.update_simulation()

            # 界面渲染
            self.screen.fill((30, 30, 30))
            self.draw_grid()
            self.draw_ui()
            pygame.display.flip()
            self.clock.tick(15)

    def update_simulation(self):
        # 批量更新所有细胞状态
        for x in range(len(self.sim.grid)):
            for y in range(len(self.sim.grid[0])):
                if self.sim.grid[x][y]:
                    self.sim.update_cell(x, y)

    def draw_ui(self):
        # 绘制控制面板和基因编辑器
        pygame.draw.rect(self.screen, (50, 50, 50), self.panel_rect)
        pygame.draw.rect(self.screen, (0, 150, 0) if self.running else (150, 0, 0), self.control_rect)
        self.panel.draw(self.screen)
        # 绘制所有按钮
        self.save_button.draw(self.screen)
        self.load_button.draw(self.screen)
        self.new_button.draw(self.screen)
        self.analyze_button.draw(self.screen)
        self.random_button.draw(self.screen)
        if self.selected_genetics:
            self.update_genetic_info()
            self.cell_count_text.draw(self.screen)
            self.genetic_info_text.draw(self.screen)
            self.total_generations_text.draw(self.screen)

    def update_genetic_info(self):
        if self.selected_genetics:
            # 使用sim.gene_stats获取同种细胞数量和平均年龄
            code_tuple = tuple(self.selected_genetics.genetic_code)
            if hasattr(self.sim, 'gene_stats') and code_tuple in self.sim.gene_stats:
                cell_count = self.sim.gene_stats[code_tuple]['count']
                genetic_lifetime = self.sim.gene_stats[code_tuple]['age'] / cell_count if cell_count > 0 else 0
            else:
                cell_count = 0
                genetic_lifetime = 0
            self.cell_count_text.set_text(f"Cells: {cell_count}")
            self.genetic_info_text.set_text(f"Genetic Info: {self.selected_genetics.encode_gene()}")
            total_generations = self.sim.total_generations
            self.total_generations_text.set_text(f"Total Generations: {total_generations}")
            # 更新基因统计面板
            sim_stats = {
                'age': genetic_lifetime,
                'count': cell_count
            }
            self.gene_stat_display.update_stats(self.selected_genetics, sim_stats)

    def toggle_fullscreen(self):
        # 切换全屏模式
        self.fullscreen = not self.fullscreen
        if self.fullscreen:
            self.screen = pygame.display.set_mode((0, 0), pygame.FULLSCREEN)
        else:
            self.screen = pygame.display.set_mode((1200, 800))

    def save_simulation(self):
        if not os.path.exists('saves'):
            os.makedirs('saves')

        if self.start_time is None:
            self.start_time = time.strftime('%Y%m%d_%H%M%S')
        generation = max(cell.lifetime for row in self.sim.grid for cell in row if cell) if any(cell for row in self.sim.grid for cell in row) else 0
        filename = f'saves/simulation_{self.start_time}_gen{generation}.pkl'

        with open(filename, 'wb') as f:
            pickle.dump(self.sim, f)

    def load_simulation(self):
        filename = tkinter.filedialog.askopenfilename(initialdir='saves', title='Select simulation file',
                                              filetypes=[('Simulation files', '*.pkl')])
        if filename:
            with open(filename, 'rb') as f:
                self.sim = pickle.load(f)

    def new_simulation(self):
        self.sim = SimulationEngine(120, 80)
        for _ in range(100):
            x, y = random.randint(0, 79), random.randint(0, 59)
            if x < len(self.sim.grid) and y < len(self.sim.grid[0]):
                self.sim.grid[x][y] = Cell()

    def analyze_genetics(self):
        # 确保gene_stats字典被正确填充
        self.sim._update_gene_stats()
        
        if not self.sim.gene_stats:
            self.show_stats_dialog()
            return

        # 获取所有基因统计数据
        gene_stats = sorted(self.sim.gene_stats.items(), key=lambda x: x[1]['count'], reverse=True)
        
        # 计算总细胞数和基因分布比例
        total_cells = sum(stats['count'] for _, stats in gene_stats)
        
        # 获取优势基因(数量最多)
        dominant_gene = gene_stats[0]
        self.selected_genetics = CellGenetics(list(dominant_gene[0]))
        
        # 显示统计对话框
        self.show_stats_dialog()

    def show_stats_dialog(self):
        if not self.sim.gene_stats:
            return

        root = tkinter.Tk()
        root.title("基因统计结果")
        root.geometry("800x600")
        
        # 计算总细胞数
        total_cells = sum(stats['count'] for stats in self.sim.gene_stats.values())
        
        # 创建主框架
        main_frame = tkinter.ttk.Frame(root)
        main_frame.pack(expand=True, fill='both', padx=10, pady=10)
        
        # 顶部信息面板
        info_frame = tkinter.ttk.Frame(main_frame)
        info_frame.pack(fill='x', pady=5)
        
        tkinter.Label(info_frame, text=f"总细胞数: {total_cells}").pack(side='left')
        tkinter.Label(info_frame, text=f"基因型数量: {len(self.sim.gene_stats)}").pack(side='left', padx=10)
        
        # 基因统计表格
        tree_frame = tkinter.ttk.Frame(main_frame)
        tree_frame.pack(expand=True, fill='both')
        
        tree = tkinter.ttk.Treeview(tree_frame, columns=('基因型', '数量', '占比%', '平均年龄', '最大年龄'), show='headings')
        tree.heading('基因型', text='基因型')
        tree.heading('数量', text='数量')
        tree.heading('占比%', text='占比%')
        tree.heading('平均年龄', text='平均年龄')
        tree.heading('最大年龄', text='最大年龄')
        
        # 设置列宽
        tree.column('基因型', width=300)
        tree.column('数量', width=80)
        tree.column('占比%', width=80)
        tree.column('平均年龄', width=80)
        tree.column('最大年龄', width=80)
        
        # 添加数据
        for gene, stats in sorted(self.sim.gene_stats.items(), key=lambda x: x[1]['count'], reverse=True):
            avg_age = stats['age'] / stats['count'] if stats['count'] > 0 else 0
            percentage = (stats['count'] / total_cells) * 100 if total_cells > 0 else 0
            
            # 根据基因参数判断生活类型
            life_type = ""
            if avg_age > 50:
                life_type = "长寿型"
            elif avg_age > 20:
                life_type = "稳定型"
            else:
                life_type = "短命型"
                
            # 获取细胞特性组合
            traits = []
            genetics = CellGenetics(list(gene))
            genes = genetics.decode_gene()
            
            if genes['autotrophic_heterotrophic']['value'] > 0:
                traits.append('自养')
            else:
                traits.append('异养')
                
            if genes['defense_offense']['value'] < 0:
                traits.append('防御')
            else:
                traits.append('进攻')
                
            if genes['static_motion']['value'] == 0:
                traits.append('静止')
            else:
                traits.append('运动')
                
            if genes['low_metabolism_high_metabolism']['value'] > 20:
                traits.append('代谢低')
            else:
                traits.append('代谢高')
                
            if genes['long_life_short_life']['value'] > 20:
                traits.append('寿命长')
            else:
                traits.append('寿命短')
                
            if genes['attract_repel']['value'] > 0:
                traits.append('吸引同类')
            else:
                traits.append('排斥同类')
                
            traits_str = '，'.join(traits)
                
            tree.insert('', 'end', values=(
                str(gene), 
                stats['count'], 
                f"{percentage:.2f}",
                f"{avg_age:.2f}",
                stats.get('max_age', 0),
                f"{life_type} ({traits_str})"
            ))
        
        # 添加滚动条
        scrollbar = tkinter.ttk.Scrollbar(tree_frame, orient='vertical', command=tree.yview)
        tree.configure(yscrollcommand=scrollbar.set)
        scrollbar.pack(side='right', fill='y')
        tree.pack(side='left', expand=True, fill='both')
        
        root.mainloop()

    def generate_random_cells(self):
        for _ in range(50):
            x, y = random.randint(0, 79), random.randint(0, 59)
            if x < len(self.sim.grid) and y < len(self.sim.grid[0]):
                self.sim.grid[x][y] = Cell()

    def handle_click(self, pos):
        if self.control_rect.collidepoint(pos):
            self.running = not self.running
        elif 910 <= pos[0] <= 1000 and 650 <= pos[1] <= 670:
            self.save_simulation()
        elif 910 <= pos[0] <= 1000 and 680 <= pos[1] <= 700:
            self.load_simulation()
        elif 910 <= pos[0] <= 1000 and 710 <= pos[1] <= 730:
            self.new_simulation()
        elif 910 <= pos[0] <= 1000 and 740 <= pos[1] <= 760:
            self.analyze_genetics()
        elif 910 <= pos[0] <= 1000 and 770 <= pos[1] <= 790:
            self.generate_random_cells()
        else:
            x, y = pos
            grid_x, grid_y = x // self.grid_size, y // self.grid_size
            if 0 <= grid_x < len(self.sim.grid) and 0 <= grid_y < len(self.sim.grid[0]):
                cell = self.sim.grid[grid_x][grid_y]
                if cell:
                    self.selected_cell = cell
                    self.selected_genetics = cell.genetics


if __name__ == "__main__":
    simulator = BioSimulator()
    simulator.run()